use anyhow::Result;
use rmcp::{RoleClient, ServiceExt, service::RunningService, transport::ConfigureCommandExt};
use serde::{Deserialize, Serialize};
use std::path::PathBuf;
use std::{collections::HashMap, process::Stdio};

use crate::config::ParamConfig;
use crate::tool::{McpToolAdapter, ToolSet};

pub struct McpClient {
    pub sessions: HashMap<String, RunningService<RoleClient, ()>>,
}

impl McpClient {
    pub async fn new(conf: &ParamConfig) -> Result<Self, Box<dyn std::error::Error>> {
        let mut sessions = HashMap::new();
        for (name, param) in conf.mcp_servers.iter() {
            tracing::info!("starting server: {} {}", name, param.to_string());
            if param.server_type == "stdio" {
                let server = McpServerTransportConfig::Stdio {
                    command: param.get_command(),
                    args: param.get_args(),
                    envs: HashMap::new(),
                    cwd: param.cwd.clone(),
                };
                let client = server.start().await?;
                sessions.insert(name.clone(), client);
            } else if param.server_type == "sse" {
                let server = McpServerTransportConfig::Sse {
                    url: param.get_url(),
                };
                let client = server.start().await?;
                sessions.insert(name.clone(), client);
            } else {
                tracing::error!("unknown server type: {}", param.server_type);
            }
        }

        Ok(Self { sessions })
    }

    pub async fn clone_tools(&self) -> Result<ToolSet> {
        let mut tools = ToolSet::default();
        for (_, session) in self.sessions.iter() {
            let tools_set = session.list_all_tools().await.unwrap();
            for tool in tools_set.into_iter() {
                tracing::debug!(
                    "tool: {} {:?} {:?}",
                    tool.name,
                    tool.description,
                    tool.input_schema
                );
                tools.add_tool(McpToolAdapter::new(tool, session.peer().clone()));
            }
        }
        Ok(tools)
    }
}

#[derive(Debug, Serialize, Deserialize, Clone)]
#[serde(tag = "protocol", rename_all = "lowercase")]
pub enum McpServerTransportConfig {
    Sse {
        url: String,
    },
    Stdio {
        command: String,
        #[serde(default)]
        args: Vec<String>,
        #[serde(default)]
        envs: HashMap<String, String>,
        #[serde(default)]
        cwd: Option<PathBuf>,
    },
}

impl McpServerTransportConfig {
    pub async fn start(&self) -> Result<RunningService<RoleClient, ()>> {
        let client = match self {
            McpServerTransportConfig::Sse { url } => {
                let transport =
                    rmcp::transport::sse_client::SseClientTransport::start(url.to_owned()).await?;
                ().serve(transport).await?
            }
            McpServerTransportConfig::Stdio {
                command, args, cwd, ..
            } => {
                tracing::debug!("starting {} {:?}", command, args);
                let transport = rmcp::transport::child_process::TokioChildProcess::new(
                    tokio::process::Command::new(command).configure(|cmd| {
                        cmd.args(args)
                            .stderr(Stdio::inherit())
                            .stdout(Stdio::inherit());
                        if let Some(cwd) = cwd {
                            cmd.current_dir(cwd);
                        }
                    }),
                )
                .unwrap();
                ().serve(transport).await?
            }
        };
        Ok(client)
    }
}
